import itertools
import json
import os
from config import PROMPT_MATRIX_CATEGORIES, OUTPUT_DIR
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class PromptMatrix:
    def __init__(self, categories=None):
        """
        初始化提示词矩阵

        参数:
            categories: 提示词类别及选项，默认为配置中的类别
        """
        self.categories = categories or PROMPT_MATRIX_CATEGORIES
        self.prompt_templates = self._load_prompt_templates()

    def _load_prompt_templates(self):
        """加载提示词模板"""
        # 简洁的中文提示词模板
        return {
            "base": "{content}",
            "with_tone": "{tone}风格{content}",
            "with_perspective": "{perspective}{content}"
        }

    def generate_matrix_combinations(self, include_categories=None):
        """
        生成提示词矩阵的所有组合

        参数:
            include_categories: 要包含的类别列表，None表示包含所有类别

        返回:
            所有可能的提示词参数组合列表
        """
        if include_categories:
            # 筛选要包含的类别
            categories = {k: v for k, v in self.categories.items() if k in include_categories}
        else:
            categories = self.categories

        # 获取所有类别的选项
        category_keys = list(categories.keys())
        category_values = [categories[key] for key in category_keys]

        # 生成所有可能的组合
        combinations = []
        for combo in itertools.product(*category_values):
            combination = dict(zip(category_keys, combo))
            combinations.append(combination)

        logger.info(f"生成了 {len(combinations)} 种提示词组合")
        return combinations

    def generate_prompt(self, base_content, parameters, template_type="base"):
        """
        根据参数生成具体提示词

        参数:
            base_content: 基础内容
            parameters: 提示词参数，如语气、详细程度等
            template_type: 模板类型

        返回:
            生成的提示词
        """
        if template_type not in self.prompt_templates:
            raise ValueError(f"未知的模板类型: {template_type}")

        template = self.prompt_templates[template_type]

        # 准备模板参数
        template_params = parameters.copy()
        template_params["content"] = base_content

        # 填充模板
        try:
            prompt = template.format(**template_params)

            # 确保提示词简洁
            prompt = prompt.replace("风格风格", "风格")  # 去重

            return prompt
        except KeyError as e:
            logger.error(f"生成提示词时缺少参数: {e}")
            return None

    def save_matrix_results(self, results, filename="prompt_matrix_results.json"):
        """保存提示词矩阵实验结果"""
        filepath = os.path.join(OUTPUT_DIR, filename)

        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                # 只保存文本信息，不包含图像数据
                text_results = []
                for result in results:
                    text_result = {
                        "parameters": result.get("parameters", {}),
                        "prompt": result.get("prompt", ""),
                        "output": result.get("output", ""),
                        "image_path": result.get("image_path", "")
                    }
                    text_results.append(text_result)

                json.dump(text_results, f, ensure_ascii=False, indent=2)
            logger.info(f"提示词矩阵结果已保存至 {filepath}")
            return True
        except Exception as e:
            logger.error(f"保存提示词矩阵结果失败: {str(e)}")
            return False